import base64
import uuid
import time
import hashlib
import json
import hmac
import random
import wave
import io

import pandas as pd
import requests
from flask import Flask, request, jsonify, session, send_from_directory, send_file
from flask_cors import CORS
import mysql.connector
import os
from urllib.parse import quote

from langchain_02_models.TTS import TTS

config = {
    'user': 'myuser',
    'password': '123456',
    'host': '127.0.0.1',
    'port': '3306',
    'database': 'dbsclab2018'
}

APP_ID = '2025418150'
APP_KEY = 'qImSlzSYZndcLisE'
URI = '/vivogpt/completions'
DOMAIN = 'api-ai.vivo.com.cn'
METHOD = 'POST'

app = Flask(__name__, static_folder="static")
CORS(app)
app.secret_key = '123456'

db_connection = mysql.connector.connect(**config)
cursor = db_connection.cursor()

base_dir = "../zht"  # PDF 文件存储路径

# 读取CSV文件并解析
def read_diseases_csv(file_path):
    encodings_to_try = ['utf-8', 'gbk', 'utf-16', 'utf-8-sig']

    for enc in encodings_to_try:
        try:
            df = pd.read_csv(file_path, encoding=enc)
            break
        except Exception as e:
            continue
    else:
        raise ValueError("无法读取CSV文件，请确认文件编码格式是否为常见编码。")

    diseases = []
    for _, row in df.iterrows():
        disease = {
            "disease": row['疾病名称'],
            "symptoms": row['主要症状'],
            "medications": row['推荐药物'],
            "temperature": round(random.uniform(36.5, 40), 1),
            "glucose": round(random.uniform(4.0, 10.0), 1),
            "lipid": round(random.uniform(2.0, 7.0), 1),
            "blood_pressure_sys": random.randint(90, 150),
            "blood_pressure_dia": random.randint(60, 100),
        }
        diseases.append(disease)

    return diseases

# 随机选择病症
def random_disease(diseases):
    return random.choice(diseases)

# 获取或创建病人
def get_or_create_patient():
    if 'patient_info' not in session:
        diseases = read_diseases_csv("../zht/diseases.csv")  # 从CSV读取疾病数据
        disease = random_disease(diseases)  # 随机选择一个病症
        mbti = random.choice(["INTJ", "ENFP", "ISTJ", "INFP", "INTP", "ENTJ", "ENTP", "INFJ", "INFP", "ENFJ", "ESFJ", "ESTJ", "ISTP", "ISFP", "ESFP"])  # 随机选择MBTI类型
        session['patient_info'] = {"disease": disease, "mbti": mbti}
        session['chat_history'] = []  # 清空历史对话
    return session['patient_info']['disease'], session['patient_info']['mbti']

# MBTI类型性格反应
MBTI_RESPONSE_STYLE = {
    "INTJ": "请告诉我具体的用药机制以及可能的副作用。",
    "INTP": "我想知道这些症状之间的逻辑关系是什么。",
    "ENTJ": "我希望尽快得到最有效的治疗方案。",
    "ENTP": "我觉得这个挺复杂的，有没有其他可能的解释？",

    "INFJ": "我有些担忧……你觉得情况会严重吗？",
    "INFP": "我有点敏感，不知道这个病是不是会有严重后果。",
    "ENFJ": "医生，我真的很想尽快恢复，我家人也很担心。",
    "ENFP": "哇，这听起来有点吓人，医生你怎么看？",

    "ISTJ": "请给我明确的诊断依据和用药建议。",
    "ISFJ": "我挺担心的，医生能再多解释一下吗？",
    "ESTJ": "我需要一个清晰的治疗计划，医生。",
    "ESFJ": "我想尽快好起来，不想让家里人担心。",

    "ISTP": "我自己查过资料，这些症状可能有多种可能。",
    "ISFP": "我有点担心副作用，但我会照你的建议试试。",
    "ESTP": "你能直接告诉我该怎么做就行。",
    "ESFP": "啊真的吗？好吧，那我先按你的说法来试试～",
}

def mbti_response(base_response, mbti_type):
    extra = MBTI_RESPONSE_STYLE.get(mbti_type, "")
    return base_response + " " + extra if extra else base_response

# 强化系统提示(每一轮都包含所有生理指标)
def build_system_prompt():
    disease, mbti = get_or_create_patient()

    temperature = disease.get('temperature', 37.0)
    glucose = disease.get('glucose', 5.0)
    lipid = disease.get('lipid', 3.0)
    blood_pressure_sys = disease.get('blood_pressure_sys', 120)
    blood_pressure_dia = disease.get('blood_pressure_dia', 80)

    background_info = (
        f"【患者隐藏设定】\n"
        f"病症: {disease['disease']}。\n"
        f"症状: {disease['symptoms']}。\n"
        f"体温: {temperature}℃。\n"
        f"血糖: {glucose} mmol/L。\n"
        f"血脂: {lipid} mmol/L。\n"
        f"血压: {blood_pressure_sys}/{blood_pressure_dia} mmHg。\n"
        f"MBTI类型: {mbti}。\n"
    )

    prompt = (
        "你现在是中国一名患者，与医生进行问诊模拟。"
        "你必须表现得像一个有血有肉的真实病人，**只在医生问及具体内容时才回答**。"
        "请不要主动说出你的疾病名称、生理指标、MBTI 类型等，除非医生明确问。"
        "你回答内容中禁止出现任何‘我是 AI’、‘我无法判断’等字眼。\n"
        "请根据医生每一个问题，真实自然地回答，行为像普通人，不要刻意表达专业名词。\n"
        + background_info
    )

    return {"role": "system", "content": prompt}

# 蓝心API签名
def gen_sign_headers(app_id, app_key, method, uri, query_params):
    timestamp = str(int(time.time()))
    nonce = uuid.uuid4().hex[:8]
    request_id = query_params.get("requestId", str(uuid.uuid4()))
    query_params["requestId"] = request_id
    canonical_query = "&".join(
        f"{quote(str(k), safe='')}={quote(str(v), safe='')}" for k, v in sorted(query_params.items())
    )
    signed_headers = (
        f"x-ai-gateway-app-id:{app_id}\n"
        f"x-ai-gateway-timestamp:{timestamp}\n"
        f"x-ai-gateway-nonce:{nonce}"
    )
    signing_str = (
        f"{method.upper()}\n"
        f"{uri}\n"
        f"{canonical_query}\n"
        f"{app_id}\n"
        f"{timestamp}\n"
        f"{signed_headers}"
    )
    signature = base64.b64encode(
        hmac.new(app_key.encode(), signing_str.encode(), hashlib.sha256).digest()
    ).decode()
    headers = {
        "Content-Type": "application/json",
        "X-AI-GATEWAY-APP-ID": app_id,
        "X-AI-GATEWAY-TIMESTAMP": timestamp,
        "X-AI-GATEWAY-NONCE": nonce,
        "X-AI-GATEWAY-SIGNATURE": signature,
        "X-AI-REQUEST-ID": request_id,
        "X-AI-GATEWAY-SIGNED-HEADERS": "x-ai-gateway-app-id;x-ai-gateway-timestamp;x-ai-gateway-nonce",
    }
    return headers, request_id

def call_bluematrix_chat(messages):
    request_id = str(uuid.uuid4())
    uri_path = "/vivogpt/completions"
    query_params = {"requestId": request_id}
    headers, _ = gen_sign_headers(APP_ID, APP_KEY, "POST", uri_path, query_params)
    url = f"https://{DOMAIN}{uri_path}?requestId={request_id}"
    response = requests.post(
        url,
        json={
            "model": "vivo-BlueLM-TB-Pro",
            "messages": messages,
            "sessionId": str(uuid.uuid4()),
            "extra": {"temperature": 0.9}
        },
        headers=headers
    )
    if response.status_code == 200:
        res_obj = response.json()
        return res_obj.get("data", {}).get("content", "Error: Invalid response")
    else:
        return f"Error: {response.status_code}, {response.text}"

def call_medical_assistant_chat(messages):
    request_id = str(uuid.uuid4())
    uri_path = "/vivogpt/completions"
    query_params = {"requestId": request_id}
    headers, _ = gen_sign_headers(APP_ID, APP_KEY, "POST", uri_path, query_params)
    url = f"https://{DOMAIN}{uri_path}?requestId={request_id}"
    response = requests.post(
        url,
        json={
            "model": "vivo-BlueLM-TB-Pro",      # 用同一个或你想用的医学助手专用的模型
            "messages": messages,
            "sessionId": str(uuid.uuid4()),
            "extra": {"temperature": 0.3}       # 医学助手一般更严谨、保守
        },
        headers=headers
    )
    if response.status_code == 200:
        res_obj = response.json()
        return res_obj.get("data", {}).get("content", "Error: Invalid response")
    else:
        return f"Error: {response.status_code}, {response.text}"


def build_assistant_system_prompt(disease):
    """医学助手专属prompt，可随意调整以使其表现不同人格或更专业"""
    background_info = (
        f"【患者信息，仅供推理】\n"
        f"症状: {disease['symptoms']}。\n"
        f"体温: {disease['temperature']}℃。\n"
        f"血糖: {disease['glucose']} mmol/L。\n"
        f"血脂: {disease['lipid']} mmol/L。\n"
        f"血压: {disease['blood_pressure_sys']}/{disease['blood_pressure_dia']} mmHg。\n"
    )
    prompt = (
        "你现在是一名医学助手，将根据患者提供的症状和体征，严谨推理可能的疾病诊断和首选药物用法。"
        "请用条理清晰、准确、医学风格的语言回答，不要隐瞒/杜撰关键指标，不要说AI字眼或本人不会诊断。"
        + background_info
    )
    return {"role": "system", "content": prompt}

def pcm2wav(pcmdata: bytes, channels=1, bits=16, sample_rate=24000):
    if bits % 8 != 0:
        raise ValueError("bits % 8 must == 0. now bits:" + str(bits))
    io_fd = io.BytesIO()
    wavfile = wave.open(io_fd, 'wb')
    wavfile.setnchannels(channels)
    wavfile.setsampwidth(bits // 8)
    wavfile.setframerate(sample_rate)
    wavfile.writeframes(pcmdata)
    wavfile.close()
    io_fd.seek(0)
    return io_fd

@app.route('/')
def index():
    return send_from_directory(app.static_folder, 'duihua.html')

@app.route('/', methods=['POST'])
def ask():
    user_message = request.form.get('user_message').strip()
    assistant_type = request.form.get('assistant_type', 'A')   # 默认为患者模式

    # 通用流程
    if user_message == "结束问答":
        if assistant_type == 'B':
            session.pop("med_assist_chat_history", None)
        else:
            session.pop("patient_info", None)
            session.pop("chat_history", None)
            session.pop('wait_diagnosis', None)
        return jsonify({"response": "本次问诊结束，下一个已生成！"})

    if session.get('wait_diagnosis', False) and assistant_type != 'B':
        return jsonify({"response": "请使用/check_diagnosis接口提交你的诊断与用药方案。"})

    if ("诊断结果" in user_message or "最终诊断" in user_message) and assistant_type != 'B':
        session['wait_diagnosis'] = True
        return jsonify({"response": "请使用/check_diagnosis接口提交本轮诊断及用药，确认后自动进入新患者。"})

    # -------------------分流处理-------------------
    disease, mbti = get_or_create_patient()
    # -------------------患者AI---------------------
    if assistant_type == 'A':
        system_message = build_system_prompt()
        chat_history = session.get("chat_history", [])
        chat_history.append({"role": "user", "content": user_message})

        resp = None
        # 硬顶指标回复
        if any(kw in user_message for kw in ["体温", "温度"]):
            base_resp = f"我量了一下，体温是{disease.get('temperature', 37.0)}℃"
            resp = mbti_response(base_resp, mbti)
        elif "血糖" in user_message:
            base_resp = f"我的血糖目前是{disease.get('glucose', 5.0)} mmol/L"
            resp = mbti_response(base_resp, mbti)
        elif "血脂" in user_message:
            base_resp = f"我的血脂目前为{disease.get('lipid', 3.0)} mmol/L"
            resp = mbti_response(base_resp, mbti)
        elif "血压" in user_message:
            base_resp = f"我的血压是{disease.get('blood_pressure_sys', 120)}/{disease.get('blood_pressure_dia', 80)} mmHg"
            resp = mbti_response(base_resp, mbti)
        elif "MBTI" in user_message.upper():
            resp = f"我是{mbti}型人格"

        if resp is not None:
            chat_history.append({"role": "assistant", "content": resp})
            session["chat_history"] = chat_history
            return jsonify({"response": resp})

        # 正常交互
        messages = [system_message] + chat_history
        response_text = call_bluematrix_chat(messages)

        for kw in ['AI医疗助手', 'AI助手', '无法提供', '没有身体']:
            if kw in response_text:
                response_text = "请问您想了解我的哪一项身体数据？"

        chat_history.append({"role": "assistant", "content": response_text})
        session["chat_history"] = chat_history
        username = session.get('CURRENT_USERNAME', 'default_username')
        cursor.execute("INSERT INTO chat_history (username, question, answer) VALUES (%s, %s, %s)",
                       (username, user_message, response_text))
        db_connection.commit()
        return jsonify({"response": response_text})

    # -------------------医学助手AI---------------------
    elif assistant_type == 'B':
        # 医学助手自己的历史，互不干扰
        med_chat_history = session.get("med_assist_chat_history", [])
        system_message = build_assistant_system_prompt(disease)
        med_chat_history.append({"role": "user", "content": user_message})
        messages = [system_message] + med_chat_history
        response_text = call_medical_assistant_chat(messages)

        for kw in ['AI医疗助手', 'AI助手', '无法提供', '没有身体', '本产品不能诊断']:
            if kw in response_text:
                response_text = "请详细描述您的症状，我来帮您分析。"

        med_chat_history.append({"role": "assistant", "content": response_text})
        session["med_assist_chat_history"] = med_chat_history
        return jsonify({"response": response_text})


# 判案后自动进新病人
@app.route("/check_diagnosis", methods=["POST"])
def check_diagnosis():
    user_diagnosis = request.form.get('diagnosis')
    disease_info = session.get('patient_info', None)

    if not disease_info:
        return jsonify({"response": "当前没有活跃的患者，请先开始问诊。"})

    correct_disease = disease_info["disease"]
    correct_med = disease_info["medications"]  # 用当前患者的推荐药物

    feedback = ""
    if correct_disease in user_diagnosis and correct_med in user_diagnosis:
        feedback = f"正确！你的最终诊断和选药都符合要求。"
    else:
        feedback = f"很遗憾，你的诊断/用药建议不完全正确。\n正确诊断应为：{correct_disease}，推荐药物：{correct_med}。"

    session.pop("patient_info", None)
    session.pop("chat_history", None)
    session['wait_diagnosis'] = False
    return jsonify({"response": feedback + " 已进入下一个患者，请继续问诊。" })

@app.route('/api/tts', methods=['POST'])
def tts_api():
    text = request.get_json().get('text', '')
    if not text.strip():
        return jsonify({'error': 'text is required'}), 400
    tts = TTS(app_id=APP_ID, app_key=APP_KEY)
    tts.open()
    pcm_buffer = tts.gen_radio(text=text)
    wav_io = pcm2wav(pcm_buffer)
    return send_file(wav_io, mimetype='audio/wav')

if __name__ == '__main__':
    app.run(debug=True, port=5000)